#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
基于向量搜索的树形结构Neo4j查询工具
功能：结合Milvus向量搜索和Neo4j图数据库进行智能查询
特点：
1. 向量相似性搜索
2. 语义理解查询
3. 混合搜索策略
4. 高精度匹配
"""

from neo4j import GraphDatabase
from pymilvus import connections, Collection
from sklearn.feature_extraction.text import TfidfVectorizer
import jieba
import re
import sys
import hashlib
import numpy as np
from collections import defaultdict


class VectorTreeQueryTool:
    def __init__(self, 
                 neo4j_uri="bolt://localhost:7687", 
                 neo4j_user="neo4j", 
                 neo4j_password="Pa@ssw0rd",
                 milvus_host="localhost",
                 milvus_port="19530",
                 milvus_collection="knowledge_lines"):
        """
        初始化向量树形查询工具
        
        Args:
            neo4j_uri: Neo4j数据库连接URI
            neo4j_user: Neo4j用户名
            neo4j_password: Neo4j密码
            milvus_host: Milvus服务器地址
            milvus_port: Milvus服务器端口
            milvus_collection: Milvus集合名称
        """
        # Neo4j连接
        try:
            self.neo4j_driver = GraphDatabase.driver(neo4j_uri, auth=(neo4j_user, neo4j_password))
            print(f"✅ 成功连接到Neo4j数据库: {neo4j_uri}")
        except Exception as e:
            print(f"❌ 连接Neo4j失败: {e}")
            self.neo4j_driver = None
        
        # Milvus连接
        try:
            connections.connect("default", host=milvus_host, port=milvus_port)
            self.milvus_collection = Collection(milvus_collection)
            self.milvus_collection.load()
            print(f"✅ 成功连接到Milvus数据库: {milvus_host}:{milvus_port}")
            print(f"✅ 加载集合: {milvus_collection}, 记录数: {self.milvus_collection.num_entities}")
        except Exception as e:
            print(f"❌ 连接Milvus失败: {e}")
            self.milvus_collection = None
        
        # 初始化向量化器（用于查询向量化）
        self.vectorizer = None
        self._init_vectorizer()
    
    def _init_vectorizer(self):
        """初始化TF-IDF向量化器，使用与存储时相同的配置"""
        try:
            # 使用与存储数据时相同的向量化器配置
            self.vectorizer = TfidfVectorizer(
                max_features=100,
                ngram_range=(1, 2),  # 与存储时保持一致
                min_df=1,
                max_df=0.95
            )

            # 尝试从Milvus获取样本数据来训练向量化器
            self._train_vectorizer_with_samples()
            print("✅ 向量化器初始化成功")
        except Exception as e:
            print(f"❌ 向量化器初始化失败: {e}")

    def _train_vectorizer_with_samples(self):
        """使用Milvus中的样本数据训练向量化器"""
        if not self.milvus_collection:
            return

        try:
            # 获取样本数据来训练向量化器
            sample_results = self.milvus_collection.query(
                expr='line_number > 0',
                limit=1000,  # 使用1000条样本训练
                output_fields=['line_text']
            )

            if sample_results:
                # 预处理样本文本
                sample_texts = []
                for result in sample_results:
                    text = result.get('line_text', '')
                    processed_text = self._preprocess_text(text)
                    if processed_text.strip():
                        sample_texts.append(processed_text)

                if sample_texts:
                    # 训练向量化器
                    self.vectorizer.fit(sample_texts)
                    self._fitted_vectorizer = True
                    print(f"✅ 向量化器训练完成 (使用{len(sample_texts)}条样本)")
                else:
                    print("⚠️  没有有效的样本文本用于训练向量化器")
            else:
                print("⚠️  无法获取样本数据来训练向量化器")

        except Exception as e:
            print(f"⚠️  向量化器训练失败: {e}")

    def _preprocess_text(self, text):
        """文本预处理，与存储时保持一致"""
        # 移除特殊字符
        text = re.sub(r'[^\w\s\u4e00-\u9fff]', ' ', text)
        text = re.sub(r'\s+', ' ', text).strip()

        # 使用jieba分词
        words = jieba.cut(text)
        text = ' '.join(words)

        # 移除多余空格
        text = re.sub(r'\s+', ' ', text).strip()
        return text
    
    def close(self):
        """关闭数据库连接"""
        if self.neo4j_driver:
            self.neo4j_driver.close()
        # Milvus连接会自动管理
    
    def extract_product_code(self, query):
        """从查询中提取产品编号"""
        product_match = re.search(r'([A-Z]{1,3}\d{3,4})', query.upper())
        return product_match.group(1) if product_match else None
    
    def vector_search_similar_content(self, query, top_k=10):
        """
        使用向量搜索查找相似内容
        
        Args:
            query: 查询文本
            top_k: 返回结果数量
            
        Returns:
            list: 相似内容列表，包含文件名、行号、文本内容和相似度分数
        """
        if not self.milvus_collection or not self.vectorizer:
            return []
        
        try:
            # 预处理查询文本，使用与存储时相同的方法
            processed_query = self._preprocess_text(query)

            # 向量化查询文本
            query_vector = self._vectorize_query_text(processed_query)

            if query_vector is None:
                return []
            
            # 执行向量搜索
            search_params = {"metric_type": "L2", "params": {"nprobe": 16}}
            
            results = self.milvus_collection.search(
                [query_vector.tolist()],
                "vector",
                search_params,
                limit=top_k,
                output_fields=["filename", "line_number", "line_text"]
            )
            
            # 处理搜索结果
            similar_contents = []
            for hits in results:
                for hit in hits:
                    similar_contents.append({
                        "filename": hit.entity.get("filename"),
                        "line_number": hit.entity.get("line_number"),
                        "line_text": hit.entity.get("line_text"),
                        "similarity_score": 1.0 / (1.0 + hit.distance),  # 转换距离为相似度
                        "distance": hit.distance
                    })
            
            return similar_contents
            
        except Exception as e:
            print(f"❌ 向量搜索失败: {e}")
            return []
    
    def _vectorize_query_text(self, text):
        """向量化查询文本"""
        try:
            # 检查向量化器是否已训练
            if not hasattr(self, '_fitted_vectorizer') or not self._fitted_vectorizer:
                print("⚠️  向量化器未训练，尝试重新训练...")
                self._train_vectorizer_with_samples()

                # 如果仍然无法训练，使用查询文本进行简单训练
                if not hasattr(self, '_fitted_vectorizer') or not self._fitted_vectorizer:
                    print("⚠️  使用查询文本进行简单训练...")
                    self.vectorizer.fit([text])
                    self._fitted_vectorizer = True

            # 向量化查询文本
            vector = self.vectorizer.transform([text]).toarray()[0]

            # 确保向量维度正确
            if len(vector) != 100:  # 默认维度
                print(f"⚠️  向量维度不匹配: {len(vector)}, 期望: 100")
                return None

            return vector

        except Exception as e:
            print(f"❌ 查询文本向量化失败: {e}")
            return None
    
    def search_tree_paths_with_vector(self, query, limit=None):
        """
        基于向量搜索的树形路径查询
        
        Args:
            query: 查询文本
            limit: 结果数量限制
            
        Returns:
            list: 查询结果列表
        """
        if not self.neo4j_driver:
            return []
        
        # 提取产品编号
        product_code = self.extract_product_code(query)
        
        print(f"🔍 向量搜索查询分析:")
        print(f"   查询文本: {query}")
        print(f"   产品编号: {product_code or '未指定'}")
        
        # 1. 首先使用向量搜索找到相似内容
        similar_contents = self.vector_search_similar_content(query, top_k=20)
        
        if not similar_contents:
            print("❌ 向量搜索未找到相似内容")
            return []
        
        print(f"✅ 向量搜索找到 {len(similar_contents)} 条相似内容")
        
        # 2. 在Neo4j中查找对应的节点
        results = []
        
        with self.neo4j_driver.session() as session:
            for content in similar_contents:
                filename = content["filename"]
                line_number = content["line_number"]
                similarity_score = content["similarity_score"]
                
                # 构建Neo4j查询条件
                params = {
                    'filename': filename,
                    'line_number': line_number
                }
                
                # 如果有产品编号，添加产品编号过滤
                where_conditions = ["n.filename = $filename", "n.line_number = $line_number"]
                if product_code:
                    where_conditions.append("n.product_code = $product_code")
                    params['product_code'] = product_code
                
                where_clause = " AND ".join(where_conditions)
                
                # 查询对应的树节点
                tree_query = f"""
                    MATCH (n:TreeNode)
                    WHERE {where_clause}
                    RETURN DISTINCT n
                """
                
                try:
                    tree_result = session.run(tree_query, params)
                    matched_nodes = [record['n'] for record in tree_result]
                    
                    for node in matched_nodes:
                        node_id = node.get('id')
                        if node_id:
                            results.append({
                                'product_code': node.get('product_code'),
                                'node_type': 'TreeNode',
                                'content': node.get('content'),
                                'filename': node.get('filename', ''),
                                'line_number': node.get('line_number', 0),
                                'keywords': node.get('keywords', []),
                                'node_id': node_id,
                                'relevance_score': similarity_score,
                                'vector_similarity': similarity_score
                            })
                
                except Exception as e:
                    print(f"❌ Neo4j查询失败: {e}")
                    continue
        
        # 3. 去重和排序
        unique_results = self._deduplicate_results(results)
        unique_results.sort(key=lambda x: x['relevance_score'], reverse=True)
        
        # 4. 应用限制
        if limit:
            unique_results = unique_results[:limit]
        
        print(f"✅ 最终找到 {len(unique_results)} 个匹配节点")
        
        # 5. 如果只有一个高相关性节点，查询其链路
        high_relevance_results = [r for r in unique_results if r['relevance_score'] >= 0.7]
        
        if len(unique_results) == 1 or len(high_relevance_results) == 1:
            target_node = unique_results[0] if len(unique_results) == 1 else high_relevance_results[0]
            print(f"🔍 找到唯一高相关性节点，正在查询其所有链路...")
            chain_results = self._get_node_chains(session, target_node)
            return chain_results
        
        return unique_results
    
    def _deduplicate_results(self, results):
        """去重结果"""
        seen_nodes = set()
        unique_results = []
        
        for result in results:
            node_id = result.get('node_id')
            if node_id and node_id not in seen_nodes:
                seen_nodes.add(node_id)
                unique_results.append(result)
        
        return unique_results

    def _get_node_chains(self, session, node):
        """获取节点的所有链路"""
        node_id = node['node_id']
        product_code = node['product_code']

        chains = []

        try:
            # 查询从根到目标节点的路径
            root_to_target_query = """
                MATCH (p:Product {code: $product_code})-[:HAS_ROOT]->(root:TreeNode)
                MATCH path = (root)-[:LEADS_TO*0..]->(target:TreeNode {id: $node_id})
                RETURN nodes(path) as root_to_target_path
                ORDER BY length(path) ASC
                LIMIT 1
            """

            root_result = session.run(root_to_target_query, {
                'product_code': product_code,
                'node_id': node_id
            })

            root_record = root_result.single()
            if not root_record:
                print("❌ 未找到从根到目标节点的路径")
                return [node]

            root_to_target_path = root_record['root_to_target_path']

            # 查询从目标节点到所有叶子节点的路径
            target_to_leaves_query = """
                MATCH (target:TreeNode {id: $node_id})
                MATCH path = (target)-[:LEADS_TO*0..]->(end)-[:RESULTS_IN]->(leaf:LeafNode)
                WHERE leaf.line_number IS NOT NULL
                RETURN DISTINCT nodes(path) as target_to_leaf_path, leaf, leaf.line_number as line_number
                ORDER BY leaf.line_number
            """

            target_to_leaves_result = session.run(target_to_leaves_query, {'node_id': node_id})

            # 收集所有记录
            all_records = list(target_to_leaves_result)
            print(f"🔍 找到 {len(all_records)} 条潜在路径")

            # 限制处理数量以提高性能
            max_process = min(len(all_records), 50)
            if len(all_records) > max_process:
                print(f"⚠️  路径数量较多，处理前{max_process}条")
                all_records = all_records[:max_process]

            # 基于完整路径内容的MD5哈希去重
            seen_path_hashes = set()

            for record in all_records:
                target_to_leaf_path = record['target_to_leaf_path']
                leaf = record['leaf']
                line_number = record['line_number']

                # 构建完整路径内容用于哈希
                full_path_content = []

                # 添加从根到目标节点的路径内容
                for path_node in root_to_target_path:
                    full_path_content.append(path_node['content'])

                # 添加从目标到叶子的路径内容（排除重复的目标节点）
                if len(target_to_leaf_path) > 1:
                    for path_node in target_to_leaf_path[1:]:
                        full_path_content.append(path_node['content'])

                # 添加叶子节点内容
                if leaf['content'] not in full_path_content:
                    full_path_content.append(leaf['content'])

                # 生成路径内容的MD5哈希
                path_content_str = " -> ".join(full_path_content)
                path_hash = hashlib.md5(path_content_str.encode('utf-8')).hexdigest()

                # 基于MD5哈希去重
                if path_hash in seen_path_hashes:
                    continue
                seen_path_hashes.add(path_hash)

                # 构建完整链路
                complete_chain = []

                # 添加从根到目标节点的路径
                for path_node in root_to_target_path:
                    complete_chain.append({
                        'content': path_node['content'],
                        'node_type': 'TreeNode',
                        'is_matched': path_node['id'] == node_id
                    })

                # 添加从目标到叶子的路径（排除重复的目标节点）
                if len(target_to_leaf_path) > 1:
                    for path_node in target_to_leaf_path[1:]:
                        complete_chain.append({
                            'content': path_node['content'],
                            'node_type': 'TreeNode',
                            'is_matched': False
                        })

                # 添加叶子节点（如果内容不重复）
                if not complete_chain or leaf['content'] != complete_chain[-1]['content']:
                    complete_chain.append({
                        'content': leaf['content'],
                        'node_type': 'LeafNode',
                        'is_matched': False
                    })

                chains.append({
                    'product_code': product_code,
                    'chain': complete_chain,
                    'matched_node': node,
                    'filename': node['filename'],
                    'line_number': node['line_number']
                })

            print(f"✅ 找到 {len(chains)} 条链路")
            return chains

        except Exception as e:
            print(f"❌ 查询链路失败: {e}")
            return [node]

    def search_by_node_name(self, node_name, product_code=None):
        """根据节点名称查询其子节点"""
        if not self.neo4j_driver:
            return []

        print(f"🔍 查询节点 '{node_name}' 的子节点")
        if product_code:
            print(f"   限定产品: {product_code}")

        results = []

        with self.neo4j_driver.session() as session:
            # 构建查询条件
            where_conditions = ["parent.content = $node_name"]
            params = {'node_name': node_name}

            if product_code:
                where_conditions.append("parent.product_code = $product_code")
                params['product_code'] = product_code

            where_clause = " AND ".join(where_conditions)

            # 查询父节点的所有子节点
            children_query = f"""
                MATCH (parent:TreeNode)
                WHERE {where_clause}
                OPTIONAL MATCH (parent)-[:LEADS_TO]->(child:TreeNode)
                OPTIONAL MATCH (parent)-[:RESULTS_IN]->(leaf:LeafNode)
                RETURN DISTINCT parent, child, leaf
            """

            try:
                children_result = session.run(children_query, params)

                parent_nodes = []
                child_nodes = []
                leaf_nodes = []

                for record in children_result:
                    parent = record['parent']
                    child = record['child']
                    leaf = record['leaf']

                    if parent and parent not in parent_nodes:
                        parent_nodes.append(parent)

                    if child:
                        child_nodes.append(child)

                    if leaf:
                        leaf_nodes.append(leaf)

                if not parent_nodes:
                    print("❌ 没有找到匹配的父节点")
                    return []

                # 去重子节点
                unique_children = []
                seen_contents = set()

                # 添加树节点子节点
                for child in child_nodes:
                    content = child.get('content')
                    if content and content not in seen_contents:
                        seen_contents.add(content)
                        unique_children.append({
                            'product_code': child.get('product_code'),
                            'node_type': 'TreeNode',
                            'content': content,
                            'filename': child.get('filename', ''),
                            'line_number': child.get('line_number', 0),
                            'keywords': child.get('keywords', []),
                            'node_id': child.get('id'),
                            'relevance_score': 1.0
                        })

                # 添加叶子节点子节点
                for leaf in leaf_nodes:
                    content = leaf.get('content')
                    if content and content not in seen_contents:
                        seen_contents.add(content)
                        unique_children.append({
                            'product_code': leaf.get('product_code'),
                            'node_type': 'LeafNode',
                            'content': content,
                            'filename': leaf.get('filename', ''),
                            'line_number': leaf.get('line_number', 0),
                            'keywords': leaf.get('keywords', []),
                            'node_id': leaf.get('id'),
                            'relevance_score': 1.0
                        })

                print(f"✅ 找到 {len(unique_children)} 个子节点")

                # 如果只有一个子节点，返回从父节点开始的全链路
                if len(unique_children) == 1:
                    print(f"🔍 只有一个子节点，正在查询全链路...")
                    parent_node = parent_nodes[0]
                    chain_results = self._get_node_chains(session, {
                        'node_id': parent_node['id'],
                        'product_code': parent_node.get('product_code'),
                        'filename': parent_node.get('filename', ''),
                        'line_number': parent_node.get('line_number', 0)
                    })
                    return chain_results

                return unique_children

            except Exception as e:
                print(f"❌ 子节点查询失败: {e}")
                return []

    def display_results(self, results):
        """显示查询结果"""
        if not results:
            print("❌ 没有找到匹配的结果")
            return

        print(f"\n🎯 找到 {len(results)} 个相关结果:")
        print("=" * 80)

        for i, result in enumerate(results, 1):
            # 检查是否是链路结果
            if 'chain' in result:
                # 显示链路
                product_code = result['product_code']
                chain = result['chain']
                filename = result['filename']
                line_number = result['line_number']

                print(f"\n📋 链路 {i} (产品: {product_code})")
                print(f"📁 来源: {filename}:{line_number}")
                print("🔗 决策链路:")

                for j, step in enumerate(chain):
                    content = step['content']
                    node_type = step['node_type']
                    is_matched = step.get('is_matched', False)

                    # 节点类型图标
                    if node_type == 'LeafNode':
                        icon = "🎯"
                        prefix = "   └─"
                    else:
                        icon = "🔍" if is_matched else "📍"
                        if j == len(chain) - 1:  # 最后一个
                            prefix = "   └─"
                        else:
                            prefix = "   ├─"

                    print(f"{prefix} {icon} {content}")

                print("-" * 60)
            else:
                # 显示普通节点
                product_code = result['product_code']
                content = result['content']
                relevance = result['relevance_score']
                filename = result['filename']
                line_number = result['line_number']
                vector_similarity = result.get('vector_similarity', 0.0)

                print(f"\n📋 结果 {i} (🌳 决策节点)")
                print(f"🏷️  产品: {product_code}")
                print(f"📄 内容: {content}")
                print(f"📁 来源: {filename}:{line_number}")
                print(f"⭐ 相关性: {relevance:.2f}")
                print(f"🔍 向量相似度: {vector_similarity:.3f}")
                print("-" * 60)

    def interactive_search(self):
        """交互式搜索"""
        print("\n🔍 基于向量的树形结构交互式查询工具")
        print("💡 支持两种查询模式:")
        print("   1. 自然语言查询: 'TD504投入量', 'PK511扩增循环数', '温度设置'")
        print("   2. 节点名称查询: 'node:节点名称' 或 'node:节点名称:产品编号'")
        print("💡 输入 'quit' 或 'exit' 退出")
        print("=" * 60)

        while True:
            try:
                query = input("\n🔍 请输入查询: ").strip()

                if query.lower() in ['quit', 'exit', 'q']:
                    print("👋 再见！")
                    break

                if not query:
                    continue

                # 检查是否是节点名称查询
                if query.startswith('node:'):
                    parts = query[5:].split(':')
                    node_name = parts[0].strip()
                    product_code = parts[1].strip() if len(parts) > 1 else None

                    print(f"\n🔄 正在查询节点: {node_name}")
                    if product_code:
                        print(f"   限定产品: {product_code}")

                    results = self.search_by_node_name(node_name, product_code)
                    self.display_results(results)
                else:
                    # 向量搜索查询
                    print(f"\n🔄 正在进行向量搜索: {query}")
                    results = self.search_tree_paths_with_vector(query)
                    self.display_results(results)

            except KeyboardInterrupt:
                print("\n👋 再见！")
                break
            except Exception as e:
                print(f"❌ 查询出错: {e}")


def main():
    """主函数"""
    print("🌳 基于向量搜索的树形结构Neo4j查询工具")
    print("=" * 50)

    # 创建查询工具
    query_tool = VectorTreeQueryTool()

    if not query_tool.neo4j_driver or not query_tool.milvus_collection:
        print("❌ 无法连接到数据库")
        return

    try:
        # 检查是否有命令行参数
        if len(sys.argv) > 1:
            # 直接查询模式
            query = ' '.join(sys.argv[1:])
            print(f"🔍 查询: {query}")
            results = query_tool.search_tree_paths_with_vector(query)
            query_tool.display_results(results)
        else:
            # 交互式模式
            query_tool.interactive_search()

    finally:
        query_tool.close()


if __name__ == "__main__":
    main()
